/**
 * @file yolov5s_detect.cc
 * @author your name (you@domain.com)
 * @brief 目标检测相关事务
 * @version 0.1
 * @date 2022-12-17
 * 
 * @copyright Copyright (c) 2022
 * 
 */
#include "yolov5s_detect.h"
#include "opencv2/opencv.hpp"
//Yolov5s模型路径
static char *yolov5_model_path = "model/names.rknn";

//对外传出的结果接口
detect_result_group_t detect_result_group;

//rknn模型输入输出参数
static rknn_context ctx;  
static rknn_input inputs[1];
static rknn_output outputs[3];
static std::vector<float> out_scales;
static std::vector<uint8_t> out_zps;
//后处理参数设置
static float vis_threshold = 0.1;
static float nms_threshold = 0.5;
static float conf_threshold = 0.3;                  //该阈值决定了模型检出目标的概率、数量

#define Min(a,b) (a<b?a:b)
cv::Mat letterBox(cv::Mat orig_img, int height, int width)
{
    static int tmp_int = 0;
    char str[100];
    cv::Mat temp, resimg;
    int img_width = orig_img.cols;
    int img_height = orig_img.rows;
    float r = Min((float)height/img_height,(float)width/img_width);
    
    int new_unpad_height = int(img_height * r);
    int new_unpad_width = int(img_width * r);

    int dh = (height - new_unpad_height);
    int dw = (width -new_unpad_width);

    float pad_h = float(dh) / 2;
    float pad_w = float(dw) / 2;

    int top = (int)(pad_h + 0.5);
    int bottom = (int)(pad_h);
    int left = (int)(pad_w + 0.5);
    int right = (int)(pad_w);
    cv::resize(orig_img,temp,cv::Size(new_unpad_width,new_unpad_height), cv::INTER_LINEAR);
    cv::copyMakeBorder(temp, resimg, top, bottom, left, right,cv::BORDER_CONSTANT,(114,114,114));
    // sprintf(str, "makeBorder_%d.jpg", tmp_int);
    // cv::imwrite(str, resimg);
    // printf(str);
    // printf("done!\n");
    // tmp_int += 1;
    return resimg;
}

void init_rknn_inoutput(){
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = 640 * 640 * 3;             //width * height * channel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    for (int i = 0; i <3; i++)
    {
        outputs[i].want_float = 0;
    }
}

unsigned char *load_model(const char *filename, int *model_size)
{
    FILE *fp = fopen(filename, "rb");
    if(fp == nullptr) {
        printf("fopen %s fail!\n", filename);
        return NULL;
    }
    fseek(fp, 0, SEEK_END);
    int model_len = ftell(fp);
    unsigned char *model = (unsigned char*)malloc(model_len);
    fseek(fp, 0, SEEK_SET);
    if(model_len != fread(model, 1, model_len, fp)) {
        printf("fread %s fail!\n", filename);
        free(model);
        return NULL;
    }
    *model_size = model_len;
    if(fp) {
        fclose(fp);
    }
    return model;
}

//打印从读入rknn中获取的信息
static void printRKNNTensor(rknn_tensor_attr *attr)
{
    printf("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
           "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f\n",
           attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
           attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
           attr->qnt_type, attr->fl, attr->zp, attr->scale);
}

//载入yolov5模型, 并打印模型信息
int init_rknn_model(){
    int ret;
    int model_data_size = 0;
    unsigned char *model_data = load_model(yolov5_model_path, &model_data_size);
  
    ret = rknn_init(&ctx, model_data, model_data_size, 0);
    if(ret != 0)
    {
        std::cout <<"rknn_inputs_set failed!" <<std::endl;
        return -1;
    }        
    
    //查询rknn output的值放入
    rknn_tensor_attr output_attrs[3];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < 3; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
                         sizeof(rknn_tensor_attr));
        printRKNNTensor(&(output_attrs[i]));
    }

    //记录后处理需要的scale 和 zero_point
    for (int i = 0; i < 3; ++i)
    {   
        out_scales.push_back(output_attrs[i].scale);
        out_zps.push_back(output_attrs[i].zp);
    }
    return ret;
}
// @brief 检测一张图上可能出现的目标
// @param img 
// @return ret: -1 表示异常，0表示正常
int detect(unsigned char *img, struct ObjectDetectBuffer *pt_object_detect_buffer){                                //Todo:不能用Mat来表示图片了, 字符指针
    int ret = 0;
    inputs[0].buf = img;
    ret = rknn_inputs_set(ctx, 1, inputs);
    if(ret != 0)
        std::cout <<"rknn_inputs_set ret:" << ret <<std::endl;
    ret = rknn_run(ctx, NULL);
    if(ret != 0)
        std::cout <<"rknn_run ret:" << ret <<std::endl;
    ret = rknn_outputs_get(ctx, 3, outputs, NULL);
    if(ret != 0)
        std::cout <<"rknn_outputs_get ret:" << ret <<std::endl;
    if(ret != 0)    return ret;

    float scale_w = (float)640 / 1920;
    float scale_h = (float)640 / 1080;

    scale_h = scale_w;
    float offset_h = 0.5*(640 - 1080 * scale_w); 

    ret = post_process((uint8_t *)outputs[0].buf, (uint8_t *)outputs[1].buf, (uint8_t *)outputs[2].buf, 640, 640,
                conf_threshold, nms_threshold, vis_threshold, scale_w, scale_h, out_zps, out_scales, &detect_result_group, offset_h);
 
    //打印检测到的目标信息
    uint8_t temp = 0;
    for(int i = 0; i < detect_result_group.count; i++){
        detect_result_t *det_result = &(detect_result_group.results[i]);
        printf("%s @ (%d %d %d %d) %f\n",
               det_result->name,
               det_result->box.left, det_result->box.top, det_result->box.right, det_result->box.bottom,
               det_result->prop);
        if(strcmp(det_result->name, "yawning") == 0){
            temp = 1;
        }
        else if(strcmp(det_result->name, "eye_closing") == 0){
            temp = 2;
        }
        else if(strcmp(det_result->name, "smoking") == 0){
            temp = 3;
        }
        else if(strcmp(det_result->name, "phone_using") == 0){
            temp = 4;
        }
        else;
    }
    pt_object_detect_buffer -> buffer[pt_object_detect_buffer -> id ++] = temp;
    if(pt_object_detect_buffer -> id == 20) pt_object_detect_buffer -> id = 0;  

    return ret;
}
